# -*- coding: utf-8 -*-
# @Time    : 2023/8/3 11:38
# @Author  : Pan
# @Software: PyCharm
# @Project : VisualFramework
# @FileName: DDPM_Swin_Seg_DL


image_size = (224, 224)
max_steps = int(6.4e5)

config = {
    "type": "TargetDiffusion",
    "diffusion": {
        "noise_steps": 100,
        "beta_start": 0.01,
        "beta_end": 0.1,
        "img_size": image_size,
    },
    "base_info": {
        "step": max_steps,
        "dot": 20,
        "save_iters": 5000,
        "pretrained": None,
        "save_path": "output/model",
        "log_dir": "output/log",
        "only_last": True
    },
    "train_dataset": {
        "type": "TargetDDPMDataset",
        "batch_size": 128,
        "shuffle": True,
        "num_workers": 4,
        # "balance": True,
        "norm": None,
        "data_root": "data",
        "data_list": "data/train_list.txt",
        "transforms": [
            {
                "type": "LoadData",
                "keys": ["img", "seg"],
                "func": "cv2"
            },
            {
                "type": "ResizeByShort",
                "keys": ["img", "seg"],
                "short": [i for i in range(224, 512, 1)],
                "inter": ["bilinear"]
            },
            {
                "type": "RandPaddingCrop",
                "keys": ["img", "seg"],
                "pad_size": image_size,
                "crop_size": image_size
            },
            {
                "type": "ToTensor",
                "keys": ["img", "seg"]
            },
            {
                "type": "Normalize",
                "keys": ["img"],
                "mean": 0.5,
                "std": 0.5
            }
        ]
    },
    "eval_dataset": {
        "type": "TargetDDPMDataset",
        "mode": "predict",
        "data_root": "data",
        "data_list": "data/val.txt",
        "batch_size": 4,
        "shuffle": True,
        "num_workers": 4,
        "generate_nums": 4,
        "img_size": image_size,
        "transforms": [
            {
                "type": "LoadData",
                "keys": ["seg"],
                "func": "cv2"
            },
            {
                "type": "ResizeByShort",
                "keys": ["seg"],
                "short": [i for i in range(224, 512, 1)],
                "inter": ["bilinear"]
            },
            {
                "type": "RandPaddingCrop",
                "keys": ["seg"],
                "pad_size": image_size,
                "crop_size": image_size
            },
            {
                "type": "ToTensor",
                "keys": ["seg"]
            }
        ]
    },
    "optimizer": {
        "type": "adam",
        "lr_scheduler": {
            "type": "WarmupCosineLR",
            "learning_rate": 0.0005,
            "total_steps": max_steps,
            "warmup_steps": 10000,
            "warmup_start_lr": 1e-7,
            "end_lr": 1e-7
        },
        "decay": None
    },
    "network": {
        "type": "diffusion",
        "network": {
            "type": "Target_SwinTransformer_tiny_patch4_window7_224",
            "decoder": {
                "type": "TargetDoubleLinear",
                "embed_dim":  768,
                "patch_dim": 96,
            }
        }
    },
    "loss": {
        "loss_list": [
            {
                "type": "L2Loss"
            }
        ],
        "loss_coef": [1]
    }
}